from langchain.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from langchain_neo4j import Neo4jGraph,GraphCypherQAChain
from pydantic import BaseModel,Field,ConfigDict
from typing import Any,Type
from dotenv import load_dotenv
from model.my_chat_model import ChatModel
import os
#定义该工具接受的输入参数格式
class Neo4jInput(BaseModel):
    question:str = Field(...,description="问题")

#定义工具类
class Neo4jTool(BaseTool):
    #允许在模型中使用非标准类型
    model_config = ConfigDict(arbitrary_types_allowed=True)
    def __init__(self,**kwargs:Any):
        super().__init__(
            name="get_neo4j_tool",
            description="用于执行 Cypher 查询来回答关于公司、股东、高管、投资、关联关系的问题,必须输入的参数是问题",
            **kwargs
        )
        #加载环境变量
        load_dotenv()
    #指定工具类模型输入的参数必须是Neo4jInput类型
    args_schema :Type[BaseModel] = Neo4jInput


    #定义工具方法
    def _run(self, question:str):
        #获取模型
        chat=ChatModel()
        llm=chat.get_line_model()

        graph=Neo4jGraph(
            url=os.getenv("NEO4J_URI"),
            username=os.getenv("NEO4J_USERNAME"),
            password=os.getenv("NEO4J_PASSWORD"),
        )
        #创建提示词模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", """
        你是一个专业的 Cypher 查询生成器，能够根据用户的自然语言问题生成正确的 Cypher 查询语句。
        请严格遵循以下数据库模式：

        节点类型与属性：
        - Company: reg_date, risk_level, reg_no, name, industry, legal_rep, registered_capital, status
        - Senior: id_no, gender, birth_date, name, position
        - Shareholder: id_no, nationality, is_controlled, name, type

        关系类型与属性：
        - (Shareholder)-[:HOLDS]->(Company)：该关系有一个属性 `ratio`（浮点数，表示持股比例）
        - (Company)-[:INVESTS_IN]->(Company)：该关系有一个属性 `amount`（整数或浮点数，表示投资金额）
        - (Shareholder)-[:RELATED_TO]->(Shareholder)：该关系有一个属性 `relation`（字符串，如"亲属"、"合作伙伴"等）
        - (Senior)-[:SERVES_AS]->(Company)：该关系有一个属性 `role`（字符串，表示职务）

        生成 Cypher 查询时请注意：
        1. 必须使用 MATCH 子句查找对应节点和关系
        2. 使用 RETURN 返回所需字段（如 h.ratio, i.amount, r.relation, a.role）
        3. 不要返回整个节点，只返回需要的属性
        4. 所有字符串比较使用双引号包围
        5. 示例：
            - 查询张三持股比例：MATCH (s:Shareholder {{name: "张三"}})-[h:HOLDS]->(c:Company {{name: "某公司"}}) RETURN h.ratio 
            - 查询李四担任职位：MATCH (s:Senior {{name: "李四"}})-[a:SERVES_AS]->(c:Company) RETURN a.role, c.name 
            - 查询股东A和股东B之间的关系:MATCH (s1:Shareholder {{name: "股东A"}})-[r:RELATED_TO ]->(s2:Shareholder {{name: "股东B"}}) RETURN r.relation
            - 查询公司A在公司B投资了多少资金：MATCH (c1:Company {{name: "公司A"}})-[i:INVESTS_IN]->(c:Company {{name:"公司B"}}) RETURN i.amount
        """),
            ("human", "{question}")
        ])

        #创建链
        chain=GraphCypherQAChain.from_llm(
            llm=llm,
            graph=graph,# 图形数据库链接
            allow_dangerous_requests=True,# 允许执行Cypyhe语言
            cypher_prompt=prompt,
            verbose=True,# 输出详细日志,调试用的，可以省略
            top_k=10,# 返回结果数量,类似于 return limit 10
            validate_cypher=True  # 验证cypher语言
        )
        #提问
        rs=chain.invoke({"query":question})
        #返回答案
        return rs["result"]

# if __name__ == "__main__":
#     tool=Neo4jTool()
#     question="请你查询一下节点Senior中名称为高管2在节点Company名称为公司73中担任什么职务"
#     answer=tool._run(question)
#     print(answer)